#!/usr/bin/env python
# -*- coding:utf-8 -*-
import os,getopt,sys
import time
from libs.config.settings import VERSION,PATHS_POCS,POCS,WORKER,CONF
from libs.utils.loader import load_string_to_module
from libs.utils.threads  import run_threads
from libs.utils.requests import patch_session
from libs.utils.printer import *
from libs.config.pocsinfo import *
from libs.utils.check import *
import libs.utils.mitmSql as mitmSql

def worker():
    if not WORKER.empty():
        arg,poc = WORKER.get()
        try:
            if CONF["method"] == "verify":
                ret = poc.verify(arg)
            elif CONF["method"] == "attack":
                ret = poc.attack(arg)
            else:
                ret = poc.verify(arg)
        except Exception as e:
            ret = None
            print(e)
        if ret:
            if "text" in ret.keys():
                info("结果: " +ret["text"])
            else:
                info("结果: " + ret)
        else:
            warn("结果: 空!")

def load_pocs_code(pocs):
    #根据路径加载poc
    global POCS
    if POCS:
        POCS = []
    for poc in pocs:
        with open(poc,"r",encoding="utf-8") as f:
            model = load_string_to_module(f.read())
            POCS.append(model)
  
def spec_model_init(name:str):
    info("目标地址:{}".format(CONF['url']))
    confPocs = pocsinfo()
    pocs = confPocs.get_pocs_path(name)
    _pocs = []
    if pocs:
        for root,dirs,files in os.walk(pocs):
            files = filter(lambda x: not x.startswith("__") and x.endswith(".py"),
            files)
            _pocs.extend(map(lambda x: os.path.join(root,x),files))
    if _pocs:
        load_pocs_code(_pocs)
        start()
def spec_poc_init(name:str):
    info("目标地址:{}".format(CONF['url']))
    confPocs = pocsinfo()
    poc = confPocs.get_poc_path(name)
    _pocs = []
    if poc:
        _pocs.append(poc)
    if _pocs:
        load_pocs_code(_pocs)
        start()
def start():
    args = CONF.get("args",{})
    for poc in POCS:
        WORKER.put((CONF,poc))
    run_threads(CONF.get("thread_num",2),worker)
def end():
    info("执行完成 {0}".format(time.strftime("%X")))
def help_set():
    info('set命令参数:')
    plus("参数:url[*],url或IP地址")
    plus("参数:headers,设置头信息")
    plus("参数:thread,设置线程数,默认是2")
    plus("参数:timeout,设置响应超时时间")
def set_url(url):
    CONF['url'] = url
    info("目标地址设置成功!")

def set_headers(headers:str):
    headers = Cheaders(headers)
    if headers:
        warn("格式错误,例:k1:v1,k2:v2")
    else:
        CONF['requests']['headers'] = headers
        info("头信息设置成功.")
def set_headers_mitm(sysName:str):
    headers = mitmSql.mitm_get_headers_by_system_name("mitmhttp",sysName)
    if headers:
        CONF['requests']['headers'] = headers
        info("头信息设置成功.")
    else:
        info("头信息设置失败.")
def set_threads(num):
    CONF['thread_num'] = int(num)
    info("线程数设置成功!")
def set_timeout(timeout):
    CONF['requests']['timeout'] = int(timeout)
    info("超时时间设置成功!")
def info_pocs():
    confPocs = pocsinfo()
    pocs = confPocs.get_pocs()
    info("注册的POC模块:")
    if pocs:
        for p in pocs:
            plus(p)
def info_poc(name):
    confPocs = pocsinfo()
    pocs = confPocs.get_pocs_items(name)
    info("可执行的POC脚本:")
    if pocs:
        for p in pocs:
            fullName = name + "." + p
            pocInfo = confPocs.get_poc_info(fullName)
            plus(p + "," + pocInfo['des'])
def check_set():
    info("参数设置情况:")
    plus("URL: " + CONF['url'])
    plus("headers: " + str(CONF['requests']['headers']))
    plus("threads: " + str(CONF['thread_num']))
    plus("timeout: " + str(CONF['requests']['timeout']))
def exec_pocs(name):
    spec_model_init(name)
    end()
def exec_poc(name):
    spec_poc_init(name)
    end()

